import cv2
import numpy as np


def stereo_rectification(K1, D1, K2, D2, image_size, R, T, alpha=0.0, new_image_size=None):
    # 将所有矩阵转换为 np.float64
    K1 = K1.astype(np.float64)
    D1 = D1.astype(np.float64)
    K2 = K2.astype(np.float64)
    D2 = D2.astype(np.float64)
    R = R.astype(np.float64)
    T = T.astype(np.float64)

    R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(
        K1, D1, K2, D2,
        image_size, R, T,
        flags=cv2.CALIB_ZERO_DISPARITY,
        alpha=alpha,
        newImageSize=new_image_size
    )
    return R1, R2, P1, P2, Q, roi1, roi2


if __name__ == '__main__':
    # 示例输入参数：统一用 np.float64
    K1 = np.array([
        [3238.015711, 0.0, 1888.312299],
        [0.0, 3249.723697, 2060.253448],
        [0.0, 0.0, 1.0]
    ], dtype=np.float64)

    K2 = np.array([
        [3228.944175, 0.0, 1937.972992],
        [0.0, 3240.712264, 1964.912781],
        [0.0, 0.0, 1.0]
    ], dtype=np.float64)

    # 畸变系数，这里转换为 1x4 的二维数组
    D1 = np.array([-0.139226581, 0.087287878, 0, 0], dtype=np.float64).reshape(1, -1)
    D2 = np.array([-0.169557688, 0.148818761, 0, 0], dtype=np.float64).reshape(1, -1)

    # 图像尺寸 (宽, 高)
    image_size = (4000, 4000)

    # 左右相机之间的旋转矩阵 R (3x3)
    R = np.array([
        [0.999788038, -0.002441387, 0.02044304],
        [0.00243707, 0.999997002, 0.000236094],
        [-0.020443556, -0.000186223, 0.999790991]
    ], dtype=np.float64)

    T = np.array([-116.1710828, -1.380962318, 0.363163634], dtype=np.float64).reshape(3, 1)

    # 调用立体校正函数
    R1, R2, P1, P2, Q, roi1, roi2 = stereo_rectification(K1, D1, K2, D2, image_size, R, T, alpha=0.0)

    # 输出结果
    print("左相机校正旋转矩阵 R1:\n", R1)
    print("右相机校正旋转矩阵 R2:\n", R2)
    print("左相机投影矩阵 P1:\n", P1)
    print("右相机投影矩阵 P2:\n", P2)
    print("重投影矩阵 Q:\n", Q)
    print("左图有效区域 roi1:", roi1)
    print("右图有效区域 roi2:", roi2)
